package com.gmrz.webauthn.protocol.v1.processor.attestation;

import co.nstant.in.cbor.model.DataItem;
import com.gmrz.uaf.common.Constants;
import com.gmrz.uaf.common.GuiceUtil;
import com.gmrz.uaf.crypto.CryptoEngine;
import com.gmrz.util.Convert;
import com.gmrz.webauthn.common.CBORDecodeUtil;
import com.gmrz.webauthn.common.WebAuthnConstants;
import com.gmrz.webauthn.protocol.v1.json.WebAuthnSchemaBuilder;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnCborException;
import com.gmrz.webauthn.protocol.v1.schema.*;
import com.gmrz.webauthn.request.WebAuthnCredentialContext;
import com.google.gson.Gson;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnErrorCode;
import com.gmrz.webauthn.protocol.v1.processor.WebAuthnException;
import java.io.IOException;
import java.nio.charset.Charset;
import java.security.NoSuchAlgorithmException;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.apache.commons.codec.binary.Base64;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;

public class AndroidKeyAttestationProcessor
        extends WebAuthnBaseAssertionProcessor
        implements IAttestationProcessor
{
    private static final Logger LOG = LogManager.getLogger(AndroidKeyAttestationProcessor.class);


    public AndroidKeyAttestationProcessor(CryptoEngine cryptoEngine){
        this.cryptoEngine = cryptoEngine;
    }


    public WebAuthnConstants.AttestationResult verify(WebAuthnCredentialContext publicKeyCredential, WebAuthenticator authenticator,boolean strictMode)
            throws WebAuthnException, WebAuthnCborException {
        if ((publicKeyCredential == null) || (null == publicKeyCredential.getResponse()))
        {
            String errorMsg = "public key credential is not present for attestation verification";
            LOG.error(errorMsg);
            return WebAuthnConstants.AttestationResult.UNAVAILABLE;
        }
        List<DataItem> dataitems = CBORDecodeUtil.cborDecode(Base64.decodeBase64(publicKeyCredential.getResponse().getAttestationObject()));
        co.nstant.in.cbor.model.Map cborMap = (co.nstant.in.cbor.model.Map) dataitems.get(0);

        java.util.Map<String, Object> map = CBORDecodeUtil.cborMapConversion(cborMap);

        byte[] rawAuthData = (byte[]) map.get("authData");

        AuthenticatorData authenticatorData = GuiceUtil.getProcessorInjector().getInstance(AuthenticatorData.class);
        authenticatorData.decode(rawAuthData);


        String formatName = (String) map.get("fmt");
        AttestationFormat format = AttestationFormat.forName(formatName);

        Object attStmtObj = map.get("attStmt");

        CredentialAttestation credentialAttestation = new CredentialAttestation().withAuthenticatorData(authenticatorData).withAuthenticatorDataBytes(rawAuthData).withAttestationStatement(attStmtObj).withFormat(format);
        credentialAttestation.setAuthenticatorDataBytes(rawAuthData);

        Map<String, Object> attStmt = (Map<String, Object>) credentialAttestation.getAttestationStatement();

        String clientDataStr = publicKeyCredential.getResponse().getClientDataJSON();
        String clientDataDecodedStr = new String(Base64.decodeBase64(clientDataStr), Charset.forName(Constants.UTF_8));
        CollectedClientData clientData = WebAuthnSchemaBuilder.getGson().fromJson(clientDataDecodedStr, CollectedClientData.class);

        if (attStmt == null)
        {
            LOG.error("attStmt is missing.", WebAuthnErrorCode.ERROR_ATTESTATION_MISSING);
            return WebAuthnConstants.AttestationResult.FAILED;
        }
        LOG.info("Verifying Android Key Attestation signature using credential public key");
        try
        {
            if (attStmt.get("x5c") != null)
            {
                LOG.info("Starting basic attestation verification for Android Key Attestation format");
                List<Certificate> x509decodeCerts = extractCertChain(attStmt);
                X509Certificate attestationCertificate = (X509Certificate)x509decodeCerts.get(0);
                validateSignatureForAKA(attStmt,rawAuthData,Base64.decodeBase64(clientDataStr),attestationCertificate);
                if (strictMode) {
                    validateCertificate(x509decodeCerts, WebAuthnConstants.getWebAuthenticatorSpec(authenticator.getAaguid().toString()), true);
                }
            }
            else
            {
                LOG.error("Unsupported attestation type", WebAuthnErrorCode.STATUS_ATTESTATION_TYPE_NOT_SUPPORTED);
                return WebAuthnConstants.AttestationResult.FAILED;
            }
        }
        catch (Exception e)
        {
            LOG.error("Error in attestation verification", e);
            return WebAuthnConstants.AttestationResult.FAILED;
        }
        LOG.info("Android Key attestation verification status : " + this.attResult.name());
        return this.attResult;
    }



    private void validateSignatureForAKA(Map<String, Object> attStmt, byte[] rawAuthData,byte[] clientData,X509Certificate attestationCert)
            throws IOException, NoSuchAlgorithmException, WebAuthnException
    {
        byte[] signedData = getSignedData(rawAuthData,clientData);

        if (attestationCert == null)
        {
            String errMsg = "Failed to read attestation certificate.";
            logAndThrow(errMsg, WebAuthnErrorCode.STATUS_INVALID_ATTESTATION_DATA);
        }
        byte[] attestationSignature = (byte[])attStmt.get("sig");

        int signAlgo = getSignAlgFromAttestaionMap(attStmt);

        LOG.debug("Validating signature using alg:" + signAlgo);
        validateSignature(signAlgo, attestationCert, signedData, attestationSignature);
    }


    private List<Certificate> extractCertChain(Map<String, Object> attStmt)
            throws WebAuthnException
    {
        List<Object> x509Certs = (List)attStmt.get("x5c");
        List<Certificate> x509decodeCerts = getDecodedX509Certs(x509Certs);
        return x509decodeCerts;
    }
    private String readAndroidKeyAttestationRawData(Map<String, Object> attStmt)
            throws WebAuthnException
    {
        List<byte[]> attBytesList = (List)attStmt.get("x5c");
        if ((attBytesList == null) || (attBytesList.isEmpty()))
        {
            String errMsg = "x5c cert chain in attestation statement cannot be empty.";
            logAndThrow(errMsg, WebAuthnErrorCode.STATUS_INVALID_ATTESTATION_DATA);
        }
        List<String> certChainList = new ArrayList();
        for (byte[] certBytes : attBytesList)
        {
            String cert = Convert.toBase64(certBytes);
            certChainList.add(cert);
        }
        Gson gson = new Gson();
        String certChainListAsStr = gson.toJson(certChainList);
        return certChainListAsStr;
    }
}

